from .instance import *
import pickle
import os
from . import default
from .utils import get_default_session


class Session(object):
    def __init__(self):
        self.df_dict = {}

    def __enter__(self):
        self.__previous_default_session = get_default_session()
        default.session = self
        return self

    def __exit__(self, exc_type, exc_value, exc_tb):
        default.session = self.__previous_default_session
        if exc_tb is not None:
            return False
        else:
            return True

    def __save_to_file(self, instance, result):
        assert instance.name is not None, "You should have a name with this instance to save!"
        assert instance.manager.dir_set, "You should first set the dir of the manager!"
        filename = os.path.join(instance.manager.data_dir, instance.name + '.pkl')
        dirname = os.path.dirname(filename)
        if os.path.exists(dirname) is False:
            logging.warning('%s not exists! Now create it.' % dirname)
            os.makedirs(dirname)
        with open(filename, 'wb') as f:
            to_dump = (result, instance.info)
            pickle.dump(to_dump, f, protocol=4)

    def __load_from_file(self, instance):
        assert instance.name is not None, "You should have a name with this instance to load!"
        assert instance.manager.dir_set, "You should first set the dir of the manager!"
        filename = os.path.join(instance.manager.data_dir, instance.name + '.pkl')
        assert os.path.exists(filename), "Can't find file %s" % (filename)
        with open(filename, 'rb') as f:
            result, instance.info = pickle.load(f)
        return result

    def __ensure_list_value(self, origin_args):
        args = []
        for arg in origin_args:
            if isinstance(arg, Instance):
                if arg in self.df_dict:
                    args.append(self.df_dict[arg])
                else:
                    args.append(self.run(arg))
            elif isinstance(arg, list) or isinstance(arg, tuple):
                args.append(self.__ensure_list_value(arg))
            elif isinstance(arg, dict):
                args.append(self.__ensure_dict_value(arg))
            else:
                args.append(arg)
        return args

    def __ensure_dict_value(self, origin_kwargs):
        kwargs = {}
        for key in origin_kwargs:
            arg = origin_kwargs[key]
            if isinstance(arg, Instance):
                if arg in self.df_dict:
                    kwargs[key] = self.df_dict[arg]
                else:
                    kwargs[key] = self.run(arg)
            elif isinstance(arg, list) or isinstance(arg, tuple):
                kwargs[key] = self.__ensure_list_value(arg)
            elif isinstance(arg, dict):
                kwargs[key] = self.__ensure_dict_value(arg)
            else:
                kwargs[key] = arg
        return kwargs

    def run(self, instance):
        if isinstance(instance, LoadInstance):
            result = self.__load_from_file(instance)
        elif isinstance(instance, SaveInstance):
            result = self.__save_to_file(instance, instance.data)
        else:
            op_args = self.__ensure_list_value(instance.op_args)
            op_kwargs = self.__ensure_dict_value(instance.op_kwargs)
            result = instance.op(instance.manager, self, *op_args, **op_kwargs)
        instance.memory_saved = True
        self.df_dict[instance] = result
        return result

    def run_all(self, instance_list):
        result_list = instance_list
        for instance in instance_list:
            result_list.append(self.run(instance))
        return result_list

    def load_from_file(self, instances):
        '''
        instances: list or instance
        will save to memory
        '''
        if isinstance(instances, list):
            for instance in instance_list:
                result = self.__load_from_file(instance)
                instance.memory_saved = True
                self.df_dict[instance] = result
        else:
            instance = instances
            result = self.__load_from_file(instance)
            instance.memory_saved = True
            self.df_dict[instance] = result

    def save_to_file(self, instances):
        '''
        instances: list or instance
        '''
        if isinstance(instances, list):
            for instance in instance_list:
                if instance.memory_saved:
                    self.__save_to_file(instance, instance.df)
                else:
                    self.__save_to_file(instance, self.run(instance))
        else:
            instance = instances
            if instance.memory_saved:
                self.__save_to_file(instance, instance.df)
            else:
                self.__save_to_file(instance, self.run(instance))


default.session = Session()


def load_from_file(instances):
    sess = get_default_session()
    sess.load_from_file(instances)


def save_to_file(instances):
    sess = get_default_session()
    sess.save_to_file(instances)
